package jk

import (
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"sync"
)

func NewSSE(c Context) (*SSE, error) {
	c.Response().HeaderSet("Content-Type", "text/event-stream")
	c.Response().HeaderSet("Cache-Control", "no-cache")
	c.Response().HeaderSet("Connection", "keep-alive")
	flusher, ok := c.Response().ResponseWriter().(http.Flusher)
	if !ok {
		return nil, errors.New("ResponseWriter does not implement http.Flusher")
	}
	return newSSE(flusher, c.Response().ResponseWriter()), nil
}

type SSE struct {
	rw      sync.Mutex
	w       http.ResponseWriter
	flusher http.Flusher
}

func newSSE(flusher http.Flusher, w http.ResponseWriter) *SSE {
	return &SSE{
		rw:      sync.Mutex{},
		w:       w,
		flusher: flusher,
	}
}

func (s *SSE) Send(data string, event string, id int64, retry uint) error {
	if data == "" {
		return errors.New("data cannot be empty")
	}
	s.rw.Lock()
	defer s.rw.Unlock()
	if id > 0 {
		_, err := s.w.Write([]byte(fmt.Sprintln("id:", id)))
		if err != nil {
			return err
		}
	}
	if event != "" {
		_, err := s.w.Write([]byte(fmt.Sprintln("event:", event)))
		if err != nil {
			return err
		}
	}
	if retry > 0 {
		_, err := s.w.Write([]byte(fmt.Sprintln("retry:", retry)))
		if err != nil {
			return err
		}
	}
	_, err := s.w.Write([]byte("data: " + data + "\n\n"))
	if err != nil {
		return err
	}
	s.flusher.Flush()
	return nil
}

func (s *SSE) SendBytes(data []byte, event string, id int64, retry uint) error {
	return s.Send(string(data), event, id, retry)
}

func (s *SSE) SendData(data string) error {
	return s.Send(data, "", 0, 0)
}

func (s *SSE) SendDataBytes(data []byte) error {
	return s.Send(string(data), "", 0, 0)
}

func (s *SSE) SendDataJson(data any) error {
	if data == nil {
		return errors.New("data cannot be empty")
	}
	marshal, err := json.Marshal(data)
	if err != nil {
		return err
	}
	return s.SendDataBytes(marshal)
}
